# Testing radial basis functions.
# 23-10-26

import os
import sys
import jax
import jax.numpy as jnp
import numpy as np
import joblib
import haiku as hk
import mindspore as ms

from jax import Array, jit
from typing import Optional, Union, List
from jax.nn.initializers import lecun_normal, normal, zeros, ones
from mindspore import load_param_into_net, Tensor, nn, context

sys.path.append(os.path.dirname(sys.path[0]))

# Test for utils
# setting context & hyper parameters
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")

from cybertron.utils.rbf import LogGaussianBasis, GaussianBasis
from ms_cybertron.rbf import LogGaussianBasis as ms_LogGaussianBasis
from ms_cybertron.rbf import GaussianBasis as ms_GaussianBasis

EPSILON = 1e-2
rng = jax.random.PRNGKey(42)
print("=================Test Gaussian Basis Function=======================")
# create distance mat
dist = np.random.uniform(low=0.1, high=1.2, size=(2, 3, 3))

def gs(x, training=False):
    gs_basis = GaussianBasis()(x) # type: ignore
    return gs_basis

gs = jax.vmap(gs, in_axes=(0,))
gs_fn = hk.transform(gs, apply_rng=True)
params = gs_fn.init(rng, dist)
gs_basis = jit(gs_fn.apply)(params, rng, dist)
print("gs_basis: ", gs_basis.shape)

ms_gs = ms_GaussianBasis()
ms_dist = Tensor(dist, ms.float32)
ms_gs_basis = ms_gs(ms_dist)

print("ms_gs_basis: ", ms_gs_basis.shape)

print("[from jax] gs_basis: \n", gs_basis)
print("[from mindspore] ms_gs_basis: \n", ms_gs_basis.asnumpy())
gs_same = np.allclose(gs_basis, ms_gs_basis.asnumpy(), atol=EPSILON)
print(f"gaussian basis function is the same: {gs_same} in accuracy of {EPSILON}.")

# caculate the norm of diff
gs_diff = gs_basis - ms_gs_basis.asnumpy()
gs_norm = np.linalg.norm(gs_diff)
print(f"norm of diff: {gs_norm}")

print("=================Test Log Gaussian Basis Function====================")
# create distance mat
# dist = np.random.uniform(low=0.1, high=1.2, size=(2, 3, 3))

def lgs(x, training=False):
    lgs_basis = LogGaussianBasis(r_min=0.04)(x) # type: ignore
    return lgs_basis

lgs = jax.vmap(lgs, in_axes=(0,))
lgs_fn = hk.transform(lgs, apply_rng=True)
params = lgs_fn.init(rng, dist)
lgs_basis = jit(lgs_fn.apply)(params, rng, dist)
print("lgs_basis: ", lgs_basis.shape)

ms_lgs = ms_LogGaussianBasis()
ms_lgs_basis = ms_lgs(ms_dist)

print("ms_lgs_basis: ", ms_lgs_basis.shape)
print("[from jax] lgs_basis: \n", lgs_basis)
print("[from mindspore] lgs_basis: \n", ms_lgs_basis.asnumpy())
lgs_same = np.allclose(lgs_basis, ms_lgs_basis.asnumpy(), atol=EPSILON)
print(f"log gaussian basis function is the same: {lgs_same} in accuracy of {EPSILON}.")
# caculate the norm of diff
lgs_diff = lgs_basis - ms_lgs_basis.asnumpy()
lgs_norm = np.linalg.norm(lgs_diff)
print(f"norm of diff: {lgs_norm}")

## Another branch
print("####################################################################")
print("#####################Test on another branch#########################")
print("####################################################################")
print("=================Test Gaussian Basis Function=======================")
# create distance mat
dist = np.random.uniform(low=0.1, high=1.2, size=(2, 3, 3))

def gs(x, training=False):
    gs_basis = GaussianBasis(num_basis=96)(x) # type: ignore
    return gs_basis

gs = jax.vmap(gs, in_axes=(0,))
gs_fn = hk.transform(gs, apply_rng=True)
params = gs_fn.init(rng, dist)
gs_basis = jit(gs_fn.apply)(params, rng, dist)
print("gs_basis: ", gs_basis.shape)

ms_gs = ms_GaussianBasis(num_basis=96)
ms_dist = Tensor(dist, ms.float32)
ms_gs_basis = ms_gs(ms_dist)

print("ms_gs_basis: ", ms_gs_basis.shape)

print("[from jax] gs_basis: \n", gs_basis)
print("[from mindspore] ms_gs_basis: \n", ms_gs_basis.asnumpy())
gs_same = np.allclose(gs_basis, ms_gs_basis.asnumpy(), atol=EPSILON)
print(f"gaussian basis function is the same: {gs_same} in accuracy of {EPSILON}.")

# caculate the norm of diff
gs_diff = gs_basis - ms_gs_basis.asnumpy()
gs_norm = np.linalg.norm(gs_diff)
print(f"norm of diff: {gs_norm}")

print("=================Test Log Gaussian Basis Function====================")
# create distance mat
# dist = np.random.uniform(low=0.1, high=1.2, size=(2, 3, 3))

def lgs(x, training=False):
    lgs_basis = LogGaussianBasis(r_min=0.04, num_basis=96)(x) # type: ignore
    return lgs_basis

lgs = jax.vmap(lgs, in_axes=(0,))
lgs_fn = hk.transform(lgs, apply_rng=True)
params = lgs_fn.init(rng, dist)
lgs_basis = jit(lgs_fn.apply)(params, rng, dist)
print("lgs_basis: ", lgs_basis.shape)

ms_lgs = ms_LogGaussianBasis(num_basis=96)
ms_lgs_basis = ms_lgs(ms_dist)

print("ms_lgs_basis: ", ms_lgs_basis.shape)
print("[from jax] lgs_basis: \n", lgs_basis)
print("[from mindspore] lgs_basis: \n", ms_lgs_basis.asnumpy())
lgs_same = np.allclose(lgs_basis, ms_lgs_basis.asnumpy(), atol=EPSILON)
print(f"log gaussian basis function is the same: {lgs_same} in accuracy of {EPSILON}.")
# caculate the norm of diff
lgs_diff = lgs_basis - ms_lgs_basis.asnumpy()
lgs_norm = np.linalg.norm(lgs_diff)
print(f"norm of diff: {lgs_norm}")